"""
协同过滤
"""

import os
import numpy as np
import pandas as pd


DATA_PATH = './ml-latest-small/ratings.csv'
CACHE_DIR = './cache'


def load_data(filepath):
    cache_path = os.path.join(CACHE_DIR, 'ratings.cache')
    print('开始加载数据集...')
    if os.path.exists(cache_path):  # 判断是否存在缓存文件
        print('缓存加载中...')
        ratings = pd.read_pickle(cache_path)
        print('从缓存加载数据集完成')
    else:
        print('新数据加载中...')
        # 设置要加载字段的数据类型
        dtype = {'userId': np.int32, 'movieId': np.int32, 'rating': np.float32}
        # 加载数据前三列
        ratings = pd.read_csv(filepath, dtype=dtype, usecols=range(3))
        # 保存缓存文件
        ratings.to_pickle(cache_path)
        print('从新数据加载数据集完成')
    return ratings


class LFM(object):
    """
    隐语义模型（Latent factor model，以下简称LFM），是基于矩阵分解的推荐算法，
    在其基本算法上引入L2正则的FunkSVD算法在推荐系统领域更是广泛使用，在Spark上也有其实现。
    """
    def __init__(self, epochs, alpha, reg_p, reg_q, factors_number=10, columns=('user_id', 'item_id', 'rating')):
        """
        :param epochs: 梯度下降最高迭代次数
        :param alpha: 学习率
        :param reg_p: p矩阵正则化系数
        :param reg_q: q矩阵正则化系数
        :param factors_number: 隐因子数量
        :param columns: 字段名称
        """
        self.epochs = epochs
        self.alpha = alpha
        self.reg_p = reg_p
        self.reg_q = reg_q
        self.factors_number = factors_number
        self.columns = columns

    def fit(self, dataset):
        """
        拟合数据集
        :param dataset: ('user_id', 'item_id', 'rating')
        :return:
        """
        self.dataset = dataset
        # 用户评分数据
        self.user_ratings = dataset.groupby(self.columns[0]).agg([list])[[self.columns[1], self.columns[2]]]
        # 物品评分数据
        self.item_ratings = dataset.groupby(self.columns[1]).agg([list])[[self.columns[0], self.columns[2]]]
        # 全局平均分
        self.mean_rating = self.dataset[self.columns[2]].mean()
        # 调用sgd方法训练模型参数
        self.P, self.Q = self.sgd()

    def _init_matrix(self):
        """
        初始化P、Q矩阵，设置0-1之间的随机值作为初始值
        :return:
        """
        # User-LF
        P = dict(zip(self.user_ratings.index, np.random.rand(len(self.user_ratings), self.factors_number).astype(np.float32)))
        # Item-LF
        Q = dict(zip(self.item_ratings.index, np.random.rand(len(self.item_ratings), self.factors_number).astype(np.float32)))
        return P, Q

    def sgd(self):
        """
        利用随机梯度下降优化参数
        :return:
        """
        P, Q = self._init_matrix()
        for i in range(self.epochs):
            print('iter %d' % i)
            error_list = []
            for user_id, item_id, rating in self.dataset.itertuples(index=False):
                # 核心代码块
                v_pu = P[user_id]
                v_qi = Q[item_id]
                error = np.float32(rating - np.dot(v_pu, v_qi))
                v_pu += self.alpha * (error * v_qi - self.reg_p * v_pu)
                v_qi += self.alpha * (error * v_pu - self.reg_q * v_qi)
                P[user_id] = v_pu
                Q[item_id] = v_qi
                error_list.append(error ** 2)
            print(np.sqrt(np.mean(error_list)))
        return P, Q

    def predict(self, user_id, item_id):
        """
        预测
        :param user_id:
        :param item_id:
        :return:
        """
        return np.dot(self.P[user_id], self.Q[item_id])


if __name__ == '__main__':
    dataset = load_data(DATA_PATH)
    lfm = LFM(20, 0.02, 0.01, 0.01, columns=('userId', 'movieId', 'rating'))
    lfm.fit(dataset)
    user_id = 1
    item_id = 1
    print(lfm.predict(user_id, item_id))
